#include "tdp2_fpsystem.hpp"
#include "tdp2_string.hpp"
#include "tdp2_hash.hpp"
#include "tdp2_guard.hpp"
#include "tdp2_error.hpp"
#include "tdp2.hpp"
#include "../zlib/zlib.h"

#include "../../rdx/src/rdx/rdx_io.hpp"

using namespace RDX::Threading;
using namespace TDP::Containers;

namespace TDP
{
	namespace ResourceManagement
	{
		using namespace RDX;
		using namespace RDX::Programmability;
		using namespace RDX::ObjectManagement;
		using namespace RDX::IO;

		MFPDataStream::MFPDataStream(const Allocator &alloc)
		{
			_alloc = alloc;
			_streamActive = false;
			_f = NULL;
			_hasAborted = false;
		}

		void *MFPDataStream::zalloc(void *opaque, unsigned int items, unsigned int size)
		{
			Allocator *alloc = static_cast<Allocator *>(opaque);
			return alloc->Realloc<LargeInt>(NULL, static_cast<LargeInt>((items * size + sizeof(LargeInt) - 1) / sizeof(LargeInt)) );
		}

		void MFPDataStream::zfree(void *opaque, void *address)
		{
			Allocator *alloc = static_cast<Allocator *>(opaque);
			alloc->Free(address);
		}

		void MFPDataStream::Abort()
		{
			_hasAborted = true;
		}

		bool MFPDataStream::HasAborted() const
		{
			return _hasAborted;
		}

		bool MFPDataStream::RestartDecompression()
		{
			if(!_isCompressed)
				return true;

			if(_streamActive)
			{
				inflateEnd(&_stream);
				_streamActive = false;
			}

			memset(&_stream, 0, sizeof(z_stream));

			_stream.opaque = &_alloc;
			_stream.zfree = zfree;
			_stream.zalloc = zalloc;
			if(inflateInit2(&_stream, -MAX_WBITS) != Z_OK)
				return false;
			_streamActive = true;
			_streamEnded = false;

			_f->Seek(static_cast<LargeInt>(_dataStartLocation), IFileStream::SEEK_Start);
			_compressedDataConsumed = 0;
			_tellOffset = 0;
			_dSize = 0;
			_dOffset = 0;

			return true;
		}

		bool MFPDataStream::OpenInlineFile(RDX::IO::IFileStream *f, LargeUInt compressedSize, LargeUInt uncompressedSize, bool isCompressed)
		{
			_f = f;
			_compressedSize = compressedSize;
			_uncompressedSize = uncompressedSize;
			_isCompressed = isCompressed;
			_dataStartLocation = static_cast<LargeUInt>(f->Tell());

			return RestartDecompression();
		}

		LargeInt MFPDataStream::WriteBytes(const void *src, LargeInt numBytes)
		{
			return 0;
		}

		LargeInt MFPDataStream::ReadBytes(void *dest, LargeInt numBytes)
		{
			if(!numBytes)
				return 0;

			if(!_isCompressed)
			{
				LargeInt nRead = _f->ReadBytes(dest, numBytes);
				this->_tellOffset += nRead;
				return nRead;
			}

			if(!_streamActive)
				return 0;	// Can't do anything if the stream is dead

			UInt8 *bytesOut = static_cast<UInt8 *>(dest);
			LargeUInt nRead = 0;
			while(true)
			{
				// Read what we can from the decompressed buffer
				if(_dSize)
				{
					LargeUInt readable = _dSize;
					if(_dSize > numBytes)
						readable = numBytes;

					if(bytesOut)
					{
						memcpy(bytesOut, _decompressedData + _dOffset, readable);
						bytesOut += readable;
					}
					nRead += readable;
					_dOffset += readable;
					_dSize -= readable;
					_tellOffset += readable;
					numBytes -= readable;

					if(!numBytes)
						return nRead;	// Read everything that was requested
				}
				else
				{
					// Decompressed data was exhausted
					if(_streamEnded)
						return nRead;	// No more decompressed data

					_stream.next_out = _decompressedData;
					_stream.avail_out = DecompressedDataSize_k;

					// Try to make some input available
					if(!_stream.avail_in)
					{
						LargeUInt remainingCompressedData = _compressedSize - _compressedDataConsumed;
						LargeUInt amountFillable = remainingCompressedData;
						if(amountFillable > CompressedDataSize_k)
							amountFillable = CompressedDataSize_k;

						if(amountFillable)
						{
							LargeUInt amountAvailable = _f->ReadBytes(_compressedData, amountFillable);
							_stream.next_in = _compressedData;
							_stream.avail_in = amountAvailable;
						}
					}
					int result = inflate(&_stream, Z_SYNC_FLUSH);
					if(result == Z_STREAM_END)
						_streamEnded = true;
					else if(result != Z_OK)
						return nRead;	// Stuck

					_dSize = _stream.next_out - _decompressedData;
					_dOffset = 0;
				}
			}
		}

		void MFPDataStream::Seek(LargeInt offset, SeekType seekType)
		{
			LargeUInt uOffset;
			
			switch(seekType)
			{
			case IFileStream::SEEK_Current:
				uOffset = _tellOffset + static_cast<LargeUInt>(offset);
				break;
			case IFileStream::SEEK_End:
				uOffset = _uncompressedSize + static_cast<LargeUInt>(offset);
				break;
			case IFileStream::SEEK_Start:
				uOffset = static_cast<LargeUInt>(offset);
				break;
			}
			
			if(uOffset > _uncompressedSize)
				uOffset = _uncompressedSize;
			if(!_isCompressed)
			{
				_f->Seek(_dataStartLocation + uOffset, IFileStream::SEEK_Start);
				_tellOffset = uOffset;
				return;
			}

			if(uOffset < _tellOffset)
				RestartDecompression();					// This is earlier in the stream, so we need to restart the entire stream to get what we need
			ReadBytes(NULL, uOffset - _tellOffset);		// Decompress data until we hit the point we need
		}

		LargeInt MFPDataStream::Tell()
		{
			return static_cast<LargeInt>(_tellOffset);
		}

		MFPDataStream::~MFPDataStream()
		{
			if(_streamActive)
				inflateEnd(&_stream);
			if(_f)
				_f->Close();
		}

		void MFPDataStream::Close()
		{
			Allocator alloc = _alloc;
			this->~MFPDataStream();
			alloc.Free(this);
		}

		void MFPArchive::InitializeArchive(IFileStream *arcF)
		{
			arcF->Seek(-static_cast<long>(ZipFile::CentralDirEnd::DiskSize()), IFileStream::SEEK_End);
			ZipFile::CentralDirEnd cde;

			if(!cde.Read(arcF))
				return;

			arcF->Seek(cde.centralDirOffset, IFileStream::SEEK_Start);
			if(!_archiveFiles.TryResize(cde.centralDirTotal, ZipFile::ArchiveEntry(_alloc)))
				return;

			for(LargeUInt i=0;i<_archiveFiles.Count();i++)
			{
				if(!_archiveFiles[i].Read(arcF))
					return;
			}

			// Build actual files list
			for(LargeUInt i=0;i<_archiveFiles.Count();i++)
			{
				CString str(_alloc);
				if(!str.TrySet(_archiveFiles[i].fileName, _archiveFiles[i].fileName.Count()))
					return;

				LargeUInt *inserted = _filesByName.TryInsert(str, i, CString::TryCopy);

				if(!inserted)
					return;		// Failed to load
			}

			_loaded = true;
		}

		MFPArchive::MFPArchive(const Allocator &alloc, RDX::IO::IFileSystem *fs, const Char *path) :
			_filesByName(alloc), _archivePath(alloc)
		{
			_initialized = false;
			_alloc = alloc;
			_filesByName.SetHashFunction(CString::Hash32);
			_loaded = false;
			_baseFS = fs;

			if(!_archivePath.TrySet(path))
				return;
			_initialized = true;
		}

		RDX::IO::IFileStream *MFPArchive::OpenFileRead(const Char *path) const
		{
			if(!_loaded)
				return NULL;

			bool streamConstructed = false;
			MFPDataStream *stream = NULL;
			IFileStream *f = NULL;

			{
				CString str(_alloc);
				{
					Vector<char> translated;
					LargeUInt pathLen = DynString<Char>::CStrLen(path);

					if(!translated.TryResize(pathLen))
						return NULL;

					for(LargeUInt i=0;i<pathLen;i++)
						translated[i] = static_cast<char>(path[i]);
					if(!str.TrySet(translated, translated.Count()))
						return NULL;
				}

				const LargeUInt *pIndex = this->_filesByName.TryGetElement(str);
				if(!pIndex)
					return NULL;

				const ZipFile::ArchiveEntry &arcFile = _archiveFiles[*pIndex];

				if(arcFile.method != ZipFile::MethodDeflate_k &&
					arcFile.method != ZipFile::MethodStore_k)
					return NULL;

				stream = _alloc.CAlloc<MFPDataStream>(1);
				if(!stream)
					return NULL;

				new (stream) MFPDataStream(_alloc);
				streamConstructed = true;

				f = _baseFS->Open(_archivePath.Chars(), false);
				if(!f)
					goto failed;
				
				f->Seek(arcFile.localHeaderOffset, IFileStream::SEEK_Start);
				if(!arcFile.VerifyLocal(f))
					goto failed;
				
				if(!stream->OpenInlineFile(f, arcFile.compressedSize, arcFile.uncompressedSize, arcFile.method == ZipFile::MethodDeflate_k))
					goto failed;

				return stream;
			}

			failed:;
			{
				if(f)
				{
					f->Close();
					f = NULL;
				}
				if(streamConstructed)
				{
					stream->~MFPDataStream();
					streamConstructed = false;
				}
				if(stream)
					_alloc.Realloc<MFPDataStream>(stream, 0);
				return NULL;
			}
		}

		MFPResourceManager::MFPResourceManager(RDX::IO::IFileSystem *fs, const Allocator &alloc) :
			_archives(alloc)
		{
			_baseFS = fs;
			_alloc = alloc;
			_archives.SetHashFunction(DynString<Char>::Hash32);
		}

		RDX::IO::IFileStream *MFPResourceManager::OpenFileSFP(const Char *wpath)
		{
			Char archivePath[MAX_LOCAL_PATH];
			
			if(wcsstr(wpath, RDX_STATIC_STRING("..")) || wcsstr(wpath, RDX_STATIC_STRING("::")))
				return NULL;	// Verboten
			
			if(1 + wcslen(wpath) + wcslen(L".sfp") >= MAX_LOCAL_PATH)
				return NULL;

			wcscpy(archivePath, wpath);
			wcscat(archivePath, L".sfp");

			MFPDataStream *stream = _alloc.Realloc<MFPDataStream>(NULL, 1);
			if(!stream)
				return NULL;

			try
			{
				new (stream) MFPDataStream(_alloc);
			}
			catch(...)
			{
				stream->~MFPDataStream();
				_alloc.Realloc<MFPDataStream>(stream, 0);
				return NULL;
			}

			IFileStream *f = NULL;

			{
				f = _baseFS->Open(archivePath, false);
				if(!f)
					goto failed;

				UInt8 buf[30];

				if(f->ReadBytes(buf, 30) != 30)
					goto failed;

				UInt16 method = ZipFile::DecodeShort(buf+8);
				UInt16 fileNameSize = ZipFile::DecodeShort(buf+26);
				UInt32 compressedSize = ZipFile::DecodeLong(buf+18);
				UInt32 uncompressedSize = ZipFile::DecodeLong(buf+22);

				if(ZipFile::DecodeLong(buf) != ZipFile::CodeLocalFile_k)
					goto failed;

				f->Seek(fileNameSize, IFileStream::SEEK_Current);

				if(!stream->OpenInlineFile(f, compressedSize, uncompressedSize, method == ZipFile::MethodDeflate_k))
					goto failed;

				return stream;
			}
			
			failed:;
			{
				if(f)
				{
					f->Close();
					f = NULL;
				}
				stream->~MFPDataStream();
				_alloc.Realloc<MFPDataStream>(stream, 0);
				return NULL;
			}
		}

		// This is called from a thread and must never panic
		RDX::IO::IFileStream *MFPResourceManager::OpenFileMFP(const Char *wpath)
		{
			Char archivePath[MAX_LOCAL_PATH];

			if(wcsstr(wpath, RDX_STATIC_STRING("..")) || wcsstr(wpath, RDX_STATIC_STRING("::")))
				return NULL;	// Verboten

			// Determine the archive
			int lastSlash = -1;
			int secondLastSlash = -1;

			for(int i=0;wpath[i];i++)
			{
				if(wpath[i] == '/')
				{
					secondLastSlash = lastSlash;
					lastSlash = i;
				}
			}

			if(lastSlash != -1)
			{
				char sha256ascii[64];
				char shaPath[MAX_LOCAL_PATH];
				{
					for(size_t i=0;i<lastSlash;i++)
					{
						shaPath[i] = static_cast<char>(wpath[i]);
						if(static_cast<wchar_t>(shaPath[i]) != wpath[i])
							return NULL;	// Reject extended characters
					}
					TDP::Crypto::SHA256::SHA256State shaState;
					shaState.FeedBytes(shaPath, lastSlash);
					shaState.Terminate();
					shaState.Flush(sha256ascii);
				}
				Char sha256wide[65];
				for(int i=0;i<64;i++)
					sha256wide[i] = static_cast<Char>(sha256ascii[i]);

				size_t dirNameSpan = static_cast<size_t>(lastSlash - secondLastSlash - 1);
				Char shortPath[MAX_LOCAL_PATH];

				for(size_t i=0;i<dirNameSpan;i++)
					shortPath[i] = wpath[secondLastSlash + 1 + i];
				shortPath[dirNameSpan] = '\0';

				sha256wide[64] = static_cast<Char>(0);
				int nChars = 1 + wcslen(L"packages/_.mfp") + wcslen(shortPath) + wcslen(sha256wide);
				if(nChars >= MAX_LOCAL_PATH)
					return NULL;
				wcscpy(archivePath, L"packages/");
				wcscat(archivePath, shortPath);
				wcscat(archivePath, L"_");
				wcscat(archivePath, sha256wide);
				wcscat(archivePath, L".mfp");
			}
			else
				wcscpy(archivePath, RDX_STATIC_STRING("packages/root.mfp"));
			
			MFPArchive *arc = NULL;
			IFileStream *arcf = NULL;

			MutexLock<LightMutex> _(&_mutex);
			DynString<Char> arcPath(_alloc);
			if(!arcPath.TrySet(archivePath))
				return NULL;

			const Char *chars = arcPath.Chars();
			SingleInstance<MFPArchive> *arciPtr = _archives.TryGetElement(arcPath);

			if(arciPtr)
				 arc = arciPtr->Instance();
			else
			{
				// Going to have to create this
				arcf = _baseFS->Open(archivePath, false);

				{
					if(!arcf)
						goto failAddArchive;

					SingleInstance<MFPArchive> *arcSI = _archives.TryInsert(arcPath, SingleInstance<MFPArchive>(_alloc), DynString<Char>::TryCopy);
					if(!arcSI)
						goto failAddArchive;

					MFPArchive *mfpi = arcSI->GetAllocator().Realloc<MFPArchive>(NULL, 1);\
					if(mfpi == NULL)
						goto failAddArchive;

					new (mfpi) MFPArchive(_alloc, _baseFS, archivePath);
					if(!mfpi->InitializedOK())
						goto failAddArchive;

					arcSI->SetInstance(mfpi);

					arcSI->Instance()->InitializeArchive(arcf);

					arc = arcSI->Instance();
				}
				failAddArchive:;

				if(arcf)
					arcf->Close();
				arcf = NULL;
			}

			if(arc == NULL)
				return NULL;

			return arc->OpenFileRead(wpath);
		}

		RDX::IO::IFileStream *MFPResourceManager::OpenFileRead(const Char *path, bool mfp)
		{
			if(mfp)
				return OpenFileMFP(path);
			return OpenFileSFP(path);
		}

		// =============================================================================================
		MFPFileSystem::MFPFileSystem(IFileSystem *parentSystem, MFPResourceManager *mfpRM)
		{
			_parentSystem = parentSystem;
			_mfpRM = mfpRM;
		}

		RDX::IO::IFileStream *MFPFileSystem::Open(const Char *path, bool write)
		{
			Char subPath[MAX_LOCAL_PATH];
			if(1 + wcslen(L"rdxclasses/") + wcslen(path) >= MAX_LOCAL_PATH)
				return NULL;

			wcscpy(subPath, L"rdxclasses/");
			wcscat(subPath, path);

			if(write)
				return _parentSystem->Open(subPath, true);
			else
				return _mfpRM->OpenFileRead(subPath, false);
		}

		RDX::IO::IFileScanState *MFPFileSystem::ScanDirectory(const Char *directory, const Char *wildcard, const Char *vext)
		{
			Char subDir[MAX_LOCAL_PATH];

			if(1 + wcslen(L"rdxclasses/") + wcslen(directory) >= MAX_LOCAL_PATH)
				return NULL;
			wcscpy(subDir, L"rdxclasses/");
			wcscat(subDir, directory);

			return _parentSystem->ScanDirectory(subDir, wildcard, RDX_STATIC_STRING(".sfp"));
		}
	}
}
